from argparse import ArgumentParser
from project_utils.image_io import *
from project_utils.filtering.core import *

import numpy as np
import os.path as osp


def ordfilt_2d(
    array: np.ndarray,
    order: int,
    ks: tuple,
    sort: "callable" = np.sort,
    mode: str = "constant",
    value: int = 0,
):
    """Perform order-statistic filtering (non-linear) on the input array,
    which processes array cells based on their ranks within a window.

    `array`: a numpy ndarray in the shape of (h, w, ...).

    `order`: specifies the rank of the cell within the window.

    `ks`: size (kh, kw) of the kernel used in filtering.

    `sort`: customized sort function, which returns a sorted list
    of array cells for each row window.

    `mode`: the boundary mode specify how the array will be padded.

    `value`: a constant value that fills the padding in constant mode.

    Note that `mode` and `value` will be forwarded to function `padding_2d`.
    """

    full_array = padding_2d(array, ks, mode, value)

    def ordfilt_kernel(row_window):
        sorted_cells = sort(row_window)
        return sorted_cells[order]

    out_array = spatial_filtering(full_array, ks, ordfilt_kernel)

    ah, aw = array.shape[:2]
    oh, ow = out_array.shape[:2]

    hs = (oh - ah) // 2
    ws = (ow - aw) // 2

    filt_array = out_array[hs : hs + ah, ws : ws + aw]

    return filt_array


def medfilt_2d(
    array: np.ndarray,
    ks: tuple = (3, 3),
    sort: "callable" = np.sort,
    mode: str = "constant",
    value: int = 0,
):
    """A special case of `ordfilt_2d`, the order is calculated based on the input `ks`,
    see the `ordfilt_2d` function for the input format.
    """

    order = (ks[0] * ks[1] - 1) // 2
    filt_array = ordfilt_2d(array, order, ks, sort, mode, value)

    return filt_array


def convfilt_2d(
    array: np.ndarray,
    kernel: np.ndarray,
    size: str = "same" or "full",
    mode: str = "constant",
    value: int = 0,
):
    """Apply custom 2D convolution kernels to the input array,
    which ultilizes numpy broadcasting operation.

    `array`: a numpy ndarray in the shape of (h, w, ...).

    `kernel`: the convolution kernel in the shape of (kh, kw, ...), it's
    suggusted that the remaining demensions match those of `array`.

    `size`: size of the output result.

    `mode`: the boundary mode specify how the array will be padded.

    `value`: a constant value that fills the padding in constant mode.

    Note that `mode` and `value` will be forwarded to function `padding_2d`.
    """

    assert size in ["same", "full"] and kernel.ndim >= 2

    ah, aw = array.shape[:2]
    kh, kw = kernel.shape[:2]

    full_array = padding_2d(array, (kh, kw), mode, value)

    def convfilt_kernel(row_window: np.ndarray):
        flat_kerkel_shape = (kh * kw,) + kernel.shape[2:]
        row_kernel = kernel.reshape(*flat_kerkel_shape)[::-1]

        max_dim = max(row_window.ndim, row_kernel.ndim)

        if row_window.ndim != max_dim:
            axis = list(range(row_window.ndim, max_dim))
            row_window = np.expand_dims(row_window, axis)
        if row_kernel.ndim != max_dim:
            axis = list(range(row_kernel.ndim, max_dim))
            row_kernel = np.expand_dims(row_kernel, axis)

        return np.sum(row_window * row_kernel, axis=0)

    out_array = spatial_filtering(full_array, (kh, kw), convfilt_kernel)

    if size == "same":
        ah, aw = array.shape[:2]
        oh, ow = out_array.shape[:2]

        hs = (oh - ah) // 2
        ws = (ow - aw) // 2

        filt_array = out_array[hs : hs + ah, ws : ws + aw]

    if size == "full":
        filt_array = out_array

    return filt_array


def hsize(parameters, default_list):
    """Unify the incoming parameters into a 2D vector as the size of the convolution kernel.

    `parameters`: input parameters, positive integers and 2D vectors are legal inputs.

    `default_list`: default 2D vector associated with the fspecial method.
    """

    assert_positive_int = lambda obj: isinstance(obj, int) and obj > 0

    parameters = parameters or default_list
    if not isinstance(parameters, (list, tuple)):
        parameters = [parameters, parameters]

    assert len(parameters) == 2 and all(
        map(assert_positive_int, parameters)
    ), "Parameters should be an int or a 2D vector of ints."

    return parameters


def fspecial(type, parameters=None, sigma=None):
    """A python version of the matlab function of the same name, which supports a variety
    of predefined 2D linear spatial filters and generates filter masks.

    `type`: selection of a specific 2D linear filtering method.

    `parameters`: input parameters which related to filtering method.

    `sigma`: input additional parameter if essential.
    """

    assert type in [
        "average",
        "disk",
        "gaussian",
        "laplacian",
        "log",
        "motion",
        "prewitt",
        "sobel",
        "unsharp",
    ], "Please choose any of the above type of kernel."

    if type == "average":
        parameters = hsize(parameters, [3, 3])
        kernel = np.full(
            (parameters[0], parameters[1]), 1 / (parameters[0] * parameters[1])
        )

    elif type == "disk":
        parameters = parameters or 5
        assert parameters > 0, "Parameter should be a positive number."

        rad = parameters
        crad = int(np.ceil(rad - 0.5))
        x, y = np.meshgrid(np.arange(-crad, crad + 1), np.arange(-crad, crad + 1))
        maxxy = np.maximum(np.abs(x), np.abs(y))
        minxy = np.minimum(np.abs(x), np.abs(y))

        m1 = np.where(
            rad**2 < (maxxy + 0.5) ** 2 + (minxy - 0.5) ** 2,
            (minxy - 0.5),
            np.sqrt(np.maximum(rad**2 - (maxxy + 0.5) ** 2, 0.0)),
        )
        m2 = np.where(
            rad**2 > (maxxy - 0.5) ** 2 + (minxy + 0.5) ** 2,
            (minxy + 0.5),
            np.sqrt(np.maximum(rad**2 - (maxxy - 0.5) ** 2, 0.0)),
        )

        sgrid = (
            rad**2
            * (
                0.5 * (np.arcsin(m2 / rad) - np.arcsin(m1 / rad))
                + 0.25
                * (np.sin(2 * np.arcsin(m2 / rad)) - np.sin(2 * np.arcsin(m1 / rad)))
            )
            - (maxxy - 0.5) * (m2 - m1)
            + (m1 - minxy + 0.5)
        ) * (
            (
                (rad**2 < (maxxy + 0.5) ** 2 + (minxy + 0.5) ** 2)
                & (rad**2 > (maxxy - 0.5) ** 2 + (minxy - 0.5) ** 2)
            )
            | ((minxy == 0) & (maxxy - 0.5 < rad) & (maxxy + 0.5 >= rad))
        )
        sgrid = sgrid + ((maxxy + 0.5) ** 2 + (minxy + 0.5) ** 2 < rad**2)
        sgrid[crad, crad] = min(np.pi * rad**2, np.pi / 2)

        if crad > 0 and rad > crad - 0.5 and rad**2 < (crad - 0.5) ** 2 + 0.25:
            m1 = np.sqrt(rad**2 - (crad - 0.5) ** 2)
            m1n = m1 / rad
            sg0 = 2 * (
                rad**2
                * (
                    0.5 * np.arcsin(m1n)
                    + 0.25 * np.sin(2 * np.arcsin(m1n))
                    - m1 * (crad - 0.5)
                )
            )
            sgrid[2 * crad, crad] = sgrid[2 * crad, crad] - sg0
            sgrid[crad, 2 * crad] = sgrid[crad, 2 * crad] - sg0
            sgrid[crad, 0] = sgrid[crad, 0] - sg0
            sgrid[0, crad] = sgrid[0, crad] - sg0
        sgrid[crad, crad] = min(sgrid[crad, crad], 1)

        kernel = sgrid / np.sum(sgrid)

    elif type == "gaussian":
        parameters = hsize(parameters, [3, 3])
        sigma = sigma or 0.5
        assert sigma > 0, "Sigma should be a positive number."

        siz = [(parameter - 1) / 2 for parameter in parameters]

        x, y = np.meshgrid(
            np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1)
        )
        arg = -(x**2 + y**2) / (2 * sigma**2)

        kernel = np.exp(arg)
        kernel[kernel < np.finfo(kernel.dtype).eps * np.max(kernel)] = 0

        sumkernel = np.sum(kernel)
        if sumkernel != 0:
            kernel = kernel / sumkernel

    elif type == "laplacian":
        parameters = parameters or 0.2
        parameters = np.maximum(0, np.minimum(parameters, 1))

        h1 = parameters / (parameters + 1)
        h2 = (1 - parameters) / (parameters + 1)

        kernel = np.array([[h1, h2, h1], [h2, -4 / (parameters + 1), h2], [h1, h2, h1]])

    elif type == "log":
        parameters = hsize(parameters, [5, 5])
        sigma = sigma or 0.5
        assert sigma > 0, "Sigma should be a positive number."

        siz = [(parameter - 1) / 2 for parameter in parameters]

        x, y = np.meshgrid(
            np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1)
        )
        arg = -(x**2 + y**2) / (2 * sigma**2)

        kernel = np.exp(arg)
        kernel[kernel < np.finfo(kernel.dtype).eps * np.max(kernel)] = 0

        sumkernel = np.sum(kernel)
        if sumkernel != 0:
            kernel = kernel / sumkernel

        kernel1 = (kernel * (x**2 + y**2 - 2 * sigma**2)) / (sigma**4)
        kernel = kernel1 - np.sum(kernel1) / np.prod(parameters)

    elif type == "motion":
        parameters = parameters or 9
        sigma = sigma or 0
        assert parameters > 0, "Parameters should be a positive number."

        len = max(1, parameters)
        half = (len - 1) / 2  # rotate half length around center
        phi = (sigma % 180) / 180 * np.pi

        cosphi = np.cos(phi)
        sinphi = np.sin(phi)
        xsign = np.sign(cosphi)
        linewdt = 1

        # Define mesh for the half matrix, eps takes care of the right size for 0 & 90 rotation
        sx = int(half * cosphi + linewdt * xsign - len * np.finfo(float).eps)
        sy = int(half * sinphi + linewdt - len * np.finfo(float).eps)
        x = np.arange(0, sx + xsign, xsign)
        y = np.arange(0, sy + 1)
        x, y = np.meshgrid(x, y)

        # Define shortest distance from a pixel to the rotated line
        dist2line = y * cosphi - x * sinphi  # distance perpendicular to the line

        rad = np.sqrt(x**2 + y**2)
        # Find points beyond the line's end-point but within the line width
        lastpix = np.where((rad >= half) & (np.abs(dist2line) <= linewdt))

        # Distance to the line's end-point parallel to the line
        x2lastpix = half - np.abs((x[lastpix] + dist2line[lastpix] * sinphi) / cosphi)

        dist2line[lastpix] = np.sqrt(dist2line[lastpix] ** 2 + x2lastpix**2)
        dist2line = linewdt + np.finfo(float).eps - np.abs(dist2line)
        dist2line[dist2line < 0] = 0  # Zero out anything beyond the line width

        # Unfold half-matrix to the full size
        h = np.rot90(dist2line, 2)
        kernel = np.zeros((2 * h.shape[0] - 1, 2 * h.shape[1] - 1))
        kernel[0 : h.shape[0], 0 : h.shape[1]] = h
        kernel[
            h.shape[0] - 1 : 2 * h.shape[0], h.shape[1] - 1 : 2 * h.shape[1]
        ] = dist2line
        kernel = kernel / (np.sum(kernel) + np.finfo(float).eps * len**2)

        if cosphi > 0:
            kernel = np.flipud(kernel)

    elif type == "prewitt":
        kernel = np.array([[1, 1, 1], [0, 0, 0], [-1, -1, -1]])

    elif type == "sobel":
        kernel = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])

    elif type == "unsharp":
        parameters = parameters or 0.2
        assert parameters >= 0 and parameters <= 1, "Sigma should be in range [0,1]."

        kernel = np.array([[0, 0, 0], [0, 1, 0], [0, 0, 0]]) - fspecial(
            "laplacian", parameters
        )

    return kernel


def imfilter(
    array: np.ndarray,
    kernel: np.ndarray,
    filtering_mode: str = "corr",
    size: str = "same",
    mode: str = "constant",
    value: int = 0,
):
    """A python version of the matlab function of the same name,
    which computes convolution or correlation.

    `array`: a numpy ndarray in the shape of (h, w, ...).

    `kernel`: the convolution kernel in the shape of (kh, kw, ...), it's
    suggusted that the remaining demensions match those of `array`.

    `filtering_mode`: array or kernel arithmetic operations.

    `size`: size of the output result.

    `mode`: the boundary mode specify how the array will be padded.

    `value`: a constant value that fills the padding in constant mode.
    """

    assert filtering_mode in ["corr", "conv"]

    if filtering_mode == "corr":
        kernel_flip = kernel.reshape(kernel.size)
        kernel_flip = kernel_flip[::-1]
        kernel_flip = kernel_flip.reshape(kernel.shape)
        filt_array = convfilt_2d(array, kernel_flip, size, mode, value)

    if filtering_mode == "conv":
        filt_array = convfilt_2d(array, kernel, size, mode, value)

    return filt_array


def sobel_filter(source_image, filtering_mode):
    """Test functions for testing sobel convolution kernels.

    `source_image`: An input array describing the original image pixels.

    `filtering_mode`: array or kernel arithmetic operations.
    """

    kernel = fspecial("sobel")
    target_image = imfilter(source_image, kernel, filtering_mode)

    return target_image


def laplacian_filter(source_image, filtering_mode):
    """Test functions for testing laplacian convolution kernels.

    `source_image`: An input array describing the original image pixels.

    `filtering_mode`: array or kernel arithmetic operations.
    """

    kernel = fspecial("laplacian")
    target_image = imfilter(source_image, kernel, filtering_mode)

    return target_image


def linear_filter_test(
    source_path, target_path, type="sobel" or "laplacian", sharp=False
):
    """Test functions for testing linear filter."""

    assert type in ["sobel", "laplacian"], "Please choose among the provided filters."

    source_image = load_image_as_grayscale(source_path, "rgb")

    if type == "laplacian":
        target_image = laplacian_filter(source_image, "conv")
    if type == "sobel":
        target_image = sobel_filter(source_image, "corr")

    if sharp:
        target_image = source_image - 0.4 * target_image

    target_image = np.clip(target_image, 0, 255)
    target_image = target_image.astype(np.uint8)

    save_grayscale_image(target_path, target_image, False)


def nonlinear_filter_test(source_path, target_path):
    """Test functions for testing Nonlinear filter."""

    source_image = load_image_as_grayscale(source_path, "rgb")

    # target_image = ordfilt_2d(source_image, 7, [5, 5])
    target_image = medfilt_2d(source_image, [5, 5])

    target_image = np.clip(target_image, 0, 255)
    target_image = target_image.astype(np.uint8)

    save_grayscale_image(target_path, target_image, False)


def fspecial_test():
    """Test all types of convolution kernels appearing in the fspecial function."""

    print("==== average ====", fspecial("average", [3, 5]), sep="\n", end="\n\n")
    print("==== disk ====", fspecial("disk", 2), sep="\n", end="\n\n")
    print("==== gaussian ====", fspecial("gaussian", [3, 4], 0.4), sep="\n", end="\n\n")
    print("==== laplacian ====", fspecial("laplacian", 0.4), sep="\n", end="\n\n")
    print("==== log ====", fspecial("log", [3, 4], 0.3), sep="\n", end="\n\n")
    print("==== motion ====", fspecial("motion"), sep="\n", end="\n\n")
    print("==== prewitt ====", fspecial("prewitt"), sep="\n", end="\n\n")
    print("==== sobel ====", fspecial("sobel"), sep="\n", end="\n\n")
    print("==== unsharp ====", fspecial("unsharp", 0.4), sep="\n", end="\n\n")


def main_procedure(cmdargs):
    if cmdargs.mode == "fspecial":
        fspecial_test()

    if cmdargs.mode == "linear":
        source_path = osp.abspath(cmdargs.source)
        target_path = osp.abspath(cmdargs.target)
        linear_filter_test(source_path, target_path)

    if cmdargs.mode == "non-linear":
        source_path = osp.abspath(cmdargs.source)
        target_path = osp.abspath(cmdargs.target)
        nonlinear_filter_test(source_path, target_path)


if __name__ == "__main__":
    parser = ArgumentParser(description="Spatial Filtering Algorithms.")

    parser.add_argument(
        "mode",
        metavar="MODE",
        choices=["fspecial", "linear", "non-linear"],
        help="Select unit test: 'fspecial', 'linear', 'non-linear'.",
    )
    parser.add_argument(
        "--source", metavar="SOURCE", type=str, help="Filepath of the source image."
    )
    parser.add_argument(
        "--target", metavar="TARGET", type=str, help="Filepath of the target image."
    )

    main_procedure(parser.parse_args())
